import gradio as gr
from PIL import Image, ImageOps, ImageChops
from utils import generate_random_dir, generate_random_file
from utils.ps_canvas import PSTools
from body_seg import BodySeg, body_part
import os
from PP_Matting import PPMatting, model_names
from demo_models import pp_matting_model


def main(origin_image: Image, color: str, algorithm: str) -> Image:
    # 设置数据缓存目录
    _save_dir = generate_random_dir()
    # 数据先保存
    origin_filename = generate_random_file(dir_=_save_dir)
    origin_image.save(origin_filename)
    # 调模型
    model = pp_matting_model
    result_type = "Replace BG"
    morph_op = "Dilate"
    filename = model(image_path=origin_filename, save_dir=_save_dir,
                     result_type=result_type, bg_color=color, algorithm=algorithm)
    image = Image.open(filename)
    os.system(f"rm -rf {_save_dir}")

    return image


def demo():
    with gr.Row():
        origin_image = gr.Image(type="pil", label="原始图像")
        with gr.Column(scale=1):
            color = gr.ColorPicker(label="背景颜色", value="#d9c5c5")
            algorithm = gr.Dropdown(
                choices=model_names, value="ppmatting-512", label="选择模型")
        out_image = gr.Image(type="pil")
    gr.Button("运行").click(main, inputs=[
        origin_image, color, algorithm], outputs=[out_image])
